import torch
from torch.nn import functional as F
from  torch import nn 
from configure import *

class Yolo_conv(nn.Module):
    def __init__(self,in_channel,out_channel,kernel,stride,padding,bias=False):
        super(Yolo_conv,self).__init__()
        self.model = nn.Sequential(
            nn.Conv2d(in_channel,out_channel,kernel,stride,padding,bias=bias),
            nn.BatchNorm2d(out_channel),
            nn.LeakyReLU()
        )
    def forward(self,x):
        return self.model(x)

class Residual_conv(nn.Module):
    def __init__(self,in_channel,out_channel):
        super(Residual_conv,self).__init__()
        self.model = nn.Sequential(
            Yolo_conv(in_channel,out_channel,1,1,0),
            Yolo_conv(out_channel,in_channel,3,1,1)
        )
    def forward(self,x):
        return self.model(x)+x

class Yolo_set_conv(nn.Module):
    def __init__(self,in_channel,out_channel):
        super(Yolo_set_conv,self).__init__()
        self.model = nn.Sequential(
            Yolo_conv(in_channel,out_channel,1,1,0),
            Yolo_conv(out_channel,in_channel,3,1,1),
            Yolo_conv(in_channel,out_channel,1,1,0),
            Yolo_conv(out_channel,in_channel,3,1,1),
            Yolo_conv(in_channel,out_channel,1,1,0)
        )
    def forward(self,x):
        return self.model(x)
class Downsample(nn.Module):
    def __init__(self,in_channel,out_channel):
        super(Downsample,self).__init__()
        self.model = nn.Sequential(
            Yolo_conv(in_channel,out_channel,3,2,1)
        )
    def forward(self,x):
        return self.model(x)
class Upsample(nn.Module):
    def __init__(self):
        super(Upsample,self).__init__()
        
    def forward(self,x):
        return F.interpolate(x,scale_factor=2,mode='nearest')
    
class Yolo_all(nn.Module):
    def __init__(self):
        super(Yolo_all,self).__init__()
        self.dark52 = nn.Sequential(
            Yolo_conv(3,32,3,1,1),
            Downsample(32,64),

            Residual_conv(64,32),

            Downsample(64,128),
            
            Residual_conv(128,64),
            Residual_conv(128,64),

            Downsample(128,256),

            Residual_conv(256,128),
            Residual_conv(256,128),
            Residual_conv(256,128),
            Residual_conv(256,128),
            Residual_conv(256,128),
            Residual_conv(256,128),
            Residual_conv(256,128),
            Residual_conv(256,128)
        )
        self.dark26 = nn.Sequential(
            Downsample(256,512),

            Residual_conv(512,256),
            Residual_conv(512,256),
            Residual_conv(512,256),
            Residual_conv(512,256),
            Residual_conv(512,256),
            Residual_conv(512,256),
            Residual_conv(512,256),
            Residual_conv(512,256)
        )
        self.dark13 = nn.Sequential(
            Downsample(512,1024),

            Residual_conv(1024,512),
            Residual_conv(1024,512),
            Residual_conv(1024,512),
            Residual_conv(1024,512)
        )
        self.setconv1 = nn.Sequential(
            Yolo_set_conv(1024,512)
        )
        self.pre1 = nn.Sequential(
            Yolo_conv(512,1024,3,1,1),
            nn.Conv2d(1024,3*(5+class_num),1,1,0)
        )
        self.up13_26 = nn.Sequential(
            Yolo_conv(512,256,3,1,1),
            Upsample()
        )
        self.setconv2 = nn.Sequential(
            Yolo_set_conv(768,256)
        )
        self.pre2 = nn.Sequential(
            Yolo_conv(256,512,3,1,1),
            nn.Conv2d(512,3*(5+class_num),1,1,0)
        )
        self.up26_52 = nn.Sequential(
            Yolo_conv(256,128,3,1,1),
            Upsample()
        )
        self.setconv3 = nn.Sequential(
            Yolo_set_conv(384,128)
        )
        self.pre3 = nn.Sequential(
            Yolo_conv(128,256,3,1,1),
            nn.Conv2d(256,3*(5+class_num),1,1,0)
        )
    def forward(self,x):
        x = self.dark52(x)
        c1 = x
        x = self.dark26(x)
        c2 =x
        x = self.dark13(x)
        x = self.setconv1(x)
        pre1 = self.pre1(x)
        x =self.up13_26(x)
        x = torch.cat((c2,x),dim=1)
        x = self.setconv2(x)
        pre2 = self.pre2(x)
        x = self.up26_52(x)
        x = torch.cat((c1,x),dim=1)
        x = self.setconv3(x)
        pre3 = self.pre3(x)
        return pre1,pre2,pre3
    
if __name__ =='__main__':
    yolo = Yolo_all()
    x = torch.randn(1,3,416,416)
    print(yolo(x)[0].shape)
    print(yolo(x)[1].shape)
    print(yolo(x)[2].shape) 